from __future__ import with_statement
import sys
import scavenger
from time import sleep, time
from threading import Thread

# Check command line args.
if len(sys.argv) < 2:
    print 'Usage: batch.py imagefile -i iterations -s scheduler -t threads'
    sys.exit(1)

# Read in the image file.
with open(sys.argv[1], 'rb') as infile:
    image = infile.read()

# Check whether a number of test iterations is given.
iterations = 1
if '-i' in sys.argv:
    iterations = int(sys.argv[sys.argv.index('-i') + 1])

# Check whether a scheduler has been specified.
if '-s' in sys.argv:
    scheduler = sys.argv[sys.argv.index('-s') + 1]
else:
    print 'Usage: batch.py imagefile -i iterations -s scheduler -t threads'
    sys.exit(1)

# Check whether a number of concurrent threads has been specified.
if '-t' in sys.argv:
    threads = int(sys.argv[sys.argv.index('-t') + 1])
else:
    threads = 2

# Sleep for a little while to make sure that we discover
# the available surrogates.
sleep(2.0)

@scavenger.cprofilescavenge('len(#0)', True)
def sharpen_cprofile(image, factor):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    factor = 1.0 + float(factor)
    sharpened_image = ImageEnhance.Sharpness(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

@scavenger.cprofilescavenge('len(#0)', True)
def color_cprofile(image, factor):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    sharpened_image = ImageEnhance.Color(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

@scavenger.cprofilescavenge('len(#0)', True)
def contrast_cprofile(image, factor):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    sharpened_image = ImageEnhance.Contrast(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

@scavenger.cprofilescavenge('len(#0)', False)
def brightness_cprofile(image, factor):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    sharpened_image = ImageEnhance.Brightness(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

@scavenger.batchscavenge('len(#0)', True)
def sharpen_batch(image, factor, id=None):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    factor = 1.0 + float(factor)
    sharpened_image = ImageEnhance.Sharpness(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

@scavenger.batchscavenge('len(#0)', True)
def color_batch(image, factor, id=None):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    sharpened_image = ImageEnhance.Color(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

@scavenger.batchscavenge('len(#0)', True)
def contrast_batch(image, factor, id=None):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    sharpened_image = ImageEnhance.Contrast(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

@scavenger.batchscavenge('len(#0)', False)
def brightness_batch(image, factor, id=None):
    from PIL import Image, ImageEnhance
    from StringIO import StringIO
    sio = StringIO(image)
    pil_image = Image.open(sio)
    sharpened_image = ImageEnhance.Brightness(pil_image).enhance(factor)
    sio = StringIO()
    sharpened_image.save(sio, pil_image.format, quality=95)
    return sio.getvalue()

# Choose the function to call - based on the scheduler chosen.
function = None
if scheduler == 'cprofile':
    functions = [sharpen_cprofile, color_cprofile, contrast_cprofile, brightness_cprofile]
elif scheduler == 'batch':
    functions = [sharpen_batch, color_batch, contrast_batch, brightness_batch]
else:
    raise Exception('Unknown scheduler.')

function_names = ['sharpen', 'color', 'contrast', 'brightness']

class Worker(Thread):
    def __init__(self, functions, image, x):
        super(Worker, self).__init__()
        self._functions = functions
        self._image = image
        self._x = x
        self.happy = True

    def run(self):
        try:
            global function_names
            for x in range(0, len(self._functions)):
                self._image = self._functions[x](self._image, 0.75, id='%s%i'%(function_names[x], self._x))
        except:
            self.happy = False

# Run the test!
for x in range(0, iterations):
    start = time()
    workers = []
    try:
        for x in range(0, threads):
            w = Worker(functions, image, x+1)
            w.start()
            workers.append(w)
        for worker in workers:
            worker.join()
    except:
        print '-1',
    stop = time()

    # Check that all workers are happy.
    happy = True
    for worker in workers:
        if not worker.happy:
            happy = False
            break
    if not happy:
        print -1
    else:
        print stop - start

scavenger.shutdown()
